package com.demo.gateway;

import com.google.common.base.Optional;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import com.netflix.niws.loadbalancer.DiscoveryEnabledServer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @ClassName: GrayMetadataRule
 * @Description: 自定义负载均衡规则
 * @Author: Yu RuiXin
 * @Date: 2019/11/18 13:42
 */
public class GrayMetadataRule extends ZoneAvoidanceRule {
	public static final String META_DATA_KEY_VERSION = "service-version";

	private static final Logger logger = LoggerFactory.getLogger(GrayMetadataRule.class);

	@Override
	public Server choose(Object key) {

		List<Server> serverList = this.getPredicate().getEligibleServers(this.getLoadBalancer().getAllServers(), key);
		if (CollectionUtils.isEmpty(serverList)) {
			return null;
		}

		String serviceVersion = TokenFilter.serviceVersionTL.get();
		logger.debug("======>GrayMetadataRule:  serviceVersionTL{}", serviceVersion);

		List<Server> noMetaServerList = new ArrayList<>();
		List<Server> metaServerList = new ArrayList<>();
		for (Server server : serverList) {
			Map<String, String> metadata = ((DiscoveryEnabledServer) server).getInstanceInfo().getMetadata();
			// version策略
			String metaVersion = metadata.get(META_DATA_KEY_VERSION);
			if (!StringUtils.isEmpty(metaVersion)) {
				metaServerList.add(server);
			} else {
				noMetaServerList.add(server);
			}
		}

		if (StringUtils.isEmpty(serviceVersion)) {
			logger.debug("====> 请求未指定服务版本，将无版本号的服务进行负载均衡");
			return originChoose(noMetaServerList, key, serviceVersion);
		}else {
			Map<String, List<Server>> listMap = metaServerList.stream().collect(Collectors.groupingBy(server ->
										((DiscoveryEnabledServer) server).getInstanceInfo().getMetadata().get(META_DATA_KEY_VERSION)));
			return originChoose(listMap.get(serviceVersion), key, serviceVersion);
		}
	}

	private Server originChoose(List<Server> serverList, Object key, String serviceVersion) {
		if(CollectionUtils.isEmpty(serverList)){
			logger.error("====> 版本号:{}对应的服务列表为空，无法进行负载均衡", serviceVersion);
			return null;
		}
		Optional<Server> server = getPredicate().chooseRoundRobinAfterFiltering(serverList, key);
		if (server.isPresent()) {
			return server.get();
		} else {
			return null;
		}
	}
}